import os
import json
import hashlib
import time
import base64
import logging
from logging.handlers import RotatingFileHandler
from flask import Flask, request, jsonify, Response, stream_template
import requests
from dotenv import load_dotenv
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import struct

# 加载环境变量
load_dotenv('config.env')

# 配置日志
def setup_logging():
    """配置日志系统"""
    # 创建logs目录
    log_dir = 'logs'
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    # 配置日志格式
    log_format = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建logger
    logger = logging.getLogger('wechat_robot')
    log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
    logger.setLevel(getattr(logging, log_level))
    
    # 清除现有的handlers
    logger.handlers.clear()
    
    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(log_format)
    logger.addHandler(console_handler)
    
    # 文件处理器 - 应用日志
    max_size = int(os.getenv('LOG_MAX_SIZE', 10*1024*1024))  # 默认10MB
    backup_count = int(os.getenv('LOG_BACKUP_COUNT', 5))
    
    app_handler = RotatingFileHandler(
        os.path.join(log_dir, 'app.log'),
        maxBytes=max_size,
        backupCount=backup_count,
        encoding='utf-8'
    )
    app_handler.setLevel(logging.INFO)
    app_handler.setFormatter(log_format)
    logger.addHandler(app_handler)
    
    # 文件处理器 - 错误日志
    error_handler = RotatingFileHandler(
        os.path.join(log_dir, 'error.log'),
        maxBytes=max_size,
        backupCount=backup_count,
        encoding='utf-8'
    )
    error_handler.setLevel(logging.ERROR)
    error_handler.setFormatter(log_format)
    logger.addHandler(error_handler)
    
    return logger

# 初始化日志
logger = setup_logging()

app = Flask(__name__)

# 记录服务启动信息
logger.info("=" * 50)
logger.info("企业微信机器人服务启动")
logger.info(f"端口: {os.getenv('PORT', '53230')}")
logger.info(f"环境: {os.getenv('FLASK_ENV', 'production')}")
logger.info(f"调试模式: {os.getenv('FLASK_DEBUG', 'False')}")
logger.info("=" * 50)

class WeChatCrypto:
    """企业微信消息加解密类"""
    
    def __init__(self, token, encoding_aes_key, corp_id):
        self.token = token
        self.encoding_aes_key = encoding_aes_key
        self.corp_id = corp_id
        self.aes_key = base64.b64decode(encoding_aes_key + '=')
        logger.info(f"初始化企业微信加密模块 - Token: {token[:10]}..., CorpID: {corp_id}")
    
    def check_signature(self, signature, timestamp, nonce, msg_encrypt):
        tmp_arr = [self.token, timestamp, nonce, msg_encrypt]
        tmp_arr.sort()
        tmp_str = ''.join(tmp_arr)
        hash_str = hashlib.sha1(tmp_str.encode('utf-8')).hexdigest()
        is_valid = hash_str == signature
        logger.info(f"签名验证详情:")
        logger.info(f"  Token: {self.token}")
        logger.info(f"  Timestamp: {timestamp}")
        logger.info(f"  Nonce: {nonce}")
        logger.info(f"  MsgEncrypt: {msg_encrypt}")
        logger.info(f"  排序后字符串: {tmp_str}")
        logger.info(f"  期望签名: {signature}")
        logger.info(f"  实际签名: {hash_str}")
        logger.info(f"  验证结果: {is_valid}")
        return is_valid
    
    def decrypt(self, text):
        try:
            logger.info(f"开始解密消息: {text[:50]}...")
            encrypted_data = base64.b64decode(text)
            logger.info(f"Base64解码后长度: {len(encrypted_data)}")
            cipher = AES.new(self.aes_key, AES.MODE_CBC, self.aes_key[:16])
            decrypted = cipher.decrypt(encrypted_data)
            logger.info(f"解密后数据长度: {len(decrypted)}")
            pad_len = decrypted[-1]
            content = decrypted[:-pad_len]
            logger.info(f"去除填充后长度: {len(content)}")
            random_str = content[:16]
            msg_len = struct.unpack("!I", content[16:20])[0]
            logger.info(f"JSON内容长度(大端): {msg_len}")
            json_bytes = content[20:20+msg_len]
            corp_id_bytes = content[20+msg_len:]
            logger.info(f"提取的JSON内容: {json_bytes.decode('utf-8')}")
            expected_corp_id = corp_id_bytes.decode('utf-8')
            logger.info(f"解密得到的企业ID: {expected_corp_id}")
            logger.info(f"配置的企业ID: {self.corp_id}")
            if expected_corp_id != self.corp_id:
                logger.warning(f"企业ID不匹配: 期望={expected_corp_id}, 实际={self.corp_id}")
            result = json_bytes.decode('utf-8')
            logger.info(f"解密成功，返回内容: {result}")
            return result
        except Exception as e:
            logger.error(f"消息解密失败: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return None
    
    def encrypt(self, reply_msg, timestamp, nonce):
        try:
            json_bytes = reply_msg.encode('utf-8')
            msg_len = struct.pack("!I", len(json_bytes))
            random_str = os.urandom(16)
            text = random_str + msg_len + json_bytes + self.corp_id.encode('utf-8')
            text = pad(text, AES.block_size)
            cipher = AES.new(self.aes_key, AES.MODE_CBC, self.aes_key[:16])
            cipher_text = cipher.encrypt(text)
            base64_text = base64.b64encode(cipher_text).decode('utf-8')
            tmp_arr = [self.token, timestamp, nonce, base64_text]
            tmp_arr.sort()
            tmp_str = ''.join(tmp_arr)
            signature = hashlib.sha1(tmp_str.encode('utf-8')).hexdigest()
            logger.debug(f"消息加密成功 - 长度: {len(reply_msg)}")
            return base64_text, signature
        except Exception as e:
            logger.error(f"消息加密失败: {e}")
            return None, None

class DifyAPI:
    """Dify API 接口类"""
    
    def __init__(self, api_url, api_key):
        self.api_url = api_url
        self.api_key = api_key
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
        logger.info(f"初始化Dify API - URL: {api_url}")
    
    def send_message_stream(self, message, conversation_id=None, user_id=None):
        """发送消息到Dify并返回流式响应"""
        try:
            logger.info(f"发送流式消息到Dify - 用户: {user_id}, 消息长度: {len(message)}")
            
            # 根据Dify API文档构建请求体，使用流式模式
            payload = {
                "inputs": {
                    "accessToken": self.api_key  # 添加accessToken参数
                },
                "query": message,
                "response_mode": "streaming",  # 使用流式模式
                "user": user_id or "default_user"
            }
            
            if conversation_id:
                payload["conversation_id"] = conversation_id
            
            logger.info(f"Dify API流式请求体: {payload}")
            
            # 添加连接超时和读取超时
            response = requests.post(
                f"{self.api_url}/chat-messages",
                headers=self.headers,
                json=payload,
                timeout=(10, 60),  # (连接超时, 读取超时)
                stream=True  # 启用流式传输
            )
            
            logger.info(f"Dify API流式响应状态: {response.status_code}")
            logger.info(f"Dify API流式响应头: {dict(response.headers)}")
            
            if response.status_code == 200:
                logger.info("Dify API流式响应成功，开始处理流式数据")
                return response
            else:
                logger.error(f"Dify API 流式请求失败: {response.status_code} - {response.text}")
                return None
                
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Dify API连接失败: {e}")
            logger.error("请检查Dify API URL是否正确，以及网络连接是否正常")
            return None
        except requests.exceptions.Timeout as e:
            logger.error(f"Dify API请求超时: {e}")
            return None
        except Exception as e:
            logger.error(f"调用Dify API流式接口时发生错误: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return None

    def send_message(self, message, conversation_id=None, user_id=None):
        """发送消息到Dify（兼容模式）"""
        try:
            logger.info(f"发送消息到Dify - 用户: {user_id}, 消息长度: {len(message)}")
            
            # 根据Dify API文档构建请求体
            payload = {
                "inputs": {
                    "accessToken": self.api_key  # 添加accessToken参数
                },
                "query": message,
                "response_mode": "blocking",
                "user": user_id or "default_user"
            }
            
            if conversation_id:
                payload["conversation_id"] = conversation_id
            
            logger.info(f"Dify API请求体: {payload}")
            
            response = requests.post(
                f"{self.api_url}/chat-messages",
                headers=self.headers,
                json=payload,
                timeout=30
            )
            
            logger.info(f"Dify API响应状态: {response.status_code}")
            logger.info(f"Dify API响应头: {dict(response.headers)}")
            
            if response.status_code == 200:
                result = response.json()
                logger.info(f"Dify API完整响应: {result}")
                answer = result.get('answer', '抱歉，我暂时无法回答您的问题。')
                logger.info(f"Dify API响应成功 - 回复长度: {len(answer)}")
                return answer
            else:
                logger.error(f"Dify API 请求失败: {response.status_code} - {response.text}")
                return "抱歉，服务暂时不可用，请稍后再试。"
                
        except Exception as e:
            logger.error(f"调用Dify API时发生错误: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return "抱歉，服务暂时不可用，请稍后再试。"

# 初始化配置
wechat_crypto = WeChatCrypto(
    token=os.getenv('WECHAT_TOKEN'),
    encoding_aes_key=os.getenv('WECHAT_ENCODING_AES_KEY'),
    corp_id=os.getenv('WECHAT_CORP_ID')
)

dify_api = DifyAPI(
    api_url=os.getenv('DIFY_API_URL'),
    api_key=os.getenv('DIFY_API_KEY')
)

# ====== 新增：流式消息内存缓存 ======
from threading import Lock
stream_cache = {}
stream_cache_lock = Lock()

@app.route('/wechat', methods=['GET', 'POST'])
def wechat_webhook():
    """企业微信回调接口"""
    if request.method == 'GET':
        # 验证回调URL
        msg_signature = request.args.get('msg_signature', '')
        timestamp = request.args.get('timestamp', '')
        nonce = request.args.get('nonce', '')
        echostr = request.args.get('echostr', '')
        
        logger.info("=" * 50)
        logger.info("收到企业微信验证请求")
        logger.info(f"msg_signature: {msg_signature}")
        logger.info(f"timestamp: {timestamp}")
        logger.info(f"nonce: {nonce}")
        logger.info(f"echostr: {echostr}")
        logger.info("=" * 50)
        
        # 先使用原始echostr进行签名验证
        if wechat_crypto.check_signature(msg_signature, timestamp, nonce, echostr):
            logger.info("企业微信验证成功")
            
            # 签名验证通过后，再处理echostr
            if echostr and len(echostr) > 20:  # 如果echostr很长，可能是加密的
                logger.info("echostr可能是加密的，尝试解密...")
                try:
                    decrypted_echostr = wechat_crypto.decrypt(echostr)
                    if decrypted_echostr:
                        logger.info(f"解密后的echostr: {decrypted_echostr}")
                        return Response(decrypted_echostr, content_type='text/plain')
                    else:
                        logger.warning("echostr解密失败，返回原始值")
                        return Response(echostr, content_type='text/plain')
                except Exception as e:
                    logger.error(f"echostr解密异常: {e}")
                    return Response(echostr, content_type='text/plain')
            else:
                logger.info(f"返回原始echostr: {echostr}")
                return Response(echostr, content_type='text/plain')
        else:
            logger.warning("企业微信验证失败")
            return Response('签名验证失败', status=403, content_type='text/plain')
    
    elif request.method == 'POST':
        # 处理消息
        try:
            # 获取请求参数
            msg_signature = request.args.get('msg_signature', '')
            timestamp = request.args.get('timestamp', '')
            nonce = request.args.get('nonce', '')
            
            logger.info("=" * 50)
            logger.info("收到企业微信POST消息")
            logger.info(f"msg_signature: {msg_signature}")
            logger.info(f"timestamp: {timestamp}")
            logger.info(f"nonce: {nonce}")
            
            # 获取原始JSON数据
            json_data = request.get_json()
            if not json_data:
                logger.error("请求体不是有效的JSON格式")
                return Response('JSON格式错误', status=400, content_type='text/plain')
            
            logger.info(f"原始JSON数据: {json_data}")
            
            # 获取加密消息
            encrypt_msg = json_data.get('encrypt')
            if not encrypt_msg:
                logger.error("JSON中缺少encrypt字段")
                return Response('缺少加密字段', status=400, content_type='text/plain')
            
            logger.info(f"提取的加密内容: {encrypt_msg[:50]}...")
            
            # 验证签名
            if not wechat_crypto.check_signature(msg_signature, timestamp, nonce, encrypt_msg):
                logger.warning("消息签名验证失败")
                return Response('签名验证失败', status=403, content_type='text/plain')
            
            # 解密消息
            decrypted_json = wechat_crypto.decrypt(encrypt_msg)
            if not decrypted_json:
                logger.error("消息解密失败")
                return Response('消息解密失败', status=400, content_type='text/plain')
            
            # 解析解密后的JSON消息
            logger.info(f"解密后的JSON内容: {decrypted_json}")
            
            try:
                msg_info = json.loads(decrypted_json)
                
                # 根据实际的企业微信消息格式解析
                msg_type = msg_info.get('msgtype', '')
                from_info = msg_info.get('from', {})
                from_user = from_info.get('userid', '') if isinstance(from_info, dict) else str(from_info)
                msg_id = msg_info.get('msgid', '')
                text_info = msg_info.get('text', {})
                user_message = text_info.get('content', '') if isinstance(text_info, dict) else str(text_info)
                
                logger.info(f"解析消息成功:")
                logger.info(f"  消息类型: {msg_type}")
                logger.info(f"  发送者: {from_user}")
                logger.info(f"  消息ID: {msg_id}")
                logger.info(f"  消息内容: {user_message}")
                
            except json.JSONDecodeError as json_error:
                logger.error(f"解密后JSON解析失败: {json_error}")
                logger.error(f"解密后JSON内容: {decrypted_json}")
                return Response('JSON解析失败', status=400, content_type='text/plain')
            
            # 处理流式消息刷新请求
            if msg_type == 'stream':
                logger.info(f"收到流式消息刷新请求")
                
                # 获取流式消息ID
                stream_info = msg_info.get('stream', {})
                stream_id = stream_info.get('id', '')
                
                logger.info(f"流式消息ID: {stream_id}")
                
                # ====== 查询缓存 ======
                with stream_cache_lock:
                    cache = stream_cache.get(stream_id)
                if cache:
                    stream_message = {
                        "msgtype": "stream",
                        "stream": {
                            "id": stream_id,
                            "finish": cache.get("finish", True),
                            "content": cache.get("content", "")
                        }
                    }
                    stream_json_str = json.dumps(stream_message, ensure_ascii=False)
                    logger.info(f"构建的流式消息响应: {stream_message}")
                    encrypted_msg, signature = wechat_crypto.encrypt(stream_json_str, timestamp, nonce)
                    if encrypted_msg and signature:
                        reply_encrypted_json = {
                            "encrypt": encrypted_msg,
                            "msg_signature": signature,
                            "timestamp": timestamp,
                            "nonce": nonce
                        }
                        logger.info("返回流式消息刷新响应")
                        return jsonify(reply_encrypted_json)
                    else:
                        logger.error("加密流式消息失败")
                        return Response('流式消息加密失败', status=500, content_type='text/plain')
                else:
                    logger.warning(f"stream_id未命中缓存: {stream_id}")
                    return Response('未找到流式消息内容', status=404, content_type='text/plain')
            
            # 只处理文本消息
            elif msg_type == 'text':
                logger.info(f"处理文本消息 - 内容: {user_message}")
                
                # 生成唯一的流式消息ID
                stream_id = f"stream_{int(time.time())}_{from_user}"
                
                # 调用Dify API流式接口
                dify_stream_response = dify_api.send_message_stream(
                    message=user_message,
                    user_id=from_user
                )
                
                if dify_stream_response:
                    logger.info("开始处理Dify流式响应")
                    
                    # 存储流式消息内容（这里简化处理，实际应该存储到数据库或缓存）
                    stream_content = ""
                    processed_answers = set()  # 用于去重
                    
                    try:
                        # 处理Dify的流式数据
                        line_count = 0
                        for line in dify_stream_response.iter_lines():
                            line_count += 1
                            logger.info(f"处理第 {line_count} 行Dify流式数据")
                            
                            if line:
                                line_str = line.decode('utf-8')
                                logger.info(f"收到Dify流式数据: {line_str}")
                                
                                # 解析Dify的流式数据
                                if line_str.startswith('data: '):
                                    data_str = line_str[6:]  # 去掉 'data: ' 前缀
                                    if data_str.strip() == '[DONE]':
                                        logger.info("Dify流式响应结束")
                                        break
                                    
                                    try:
                                        data = json.loads(data_str)
                                        logger.info(f"解析Dify数据成功: {data}")
                                        
                                        # 提取Dify的回复内容
                                        answer = None
                                        
                                        # 检查不同位置的answer字段
                                        if 'answer' in data:
                                            answer = data['answer']
                                        elif 'data' in data and 'outputs' in data['data']:
                                            if 'answer' in data['data']['outputs']:
                                                answer = data['data']['outputs']['answer']
                                        
                                        if answer:
                                            # 检查是否已经处理过这个answer（去重）
                                            if answer not in processed_answers:
                                                logger.info(f"收到Dify回复片段: {answer}")
                                                stream_content += answer
                                                processed_answers.add(answer)
                                            else:
                                                logger.info(f"跳过重复的Dify回复片段")
                                            
                                        else:
                                            logger.info(f"Dify数据中没有answer字段: {data}")
                                            
                                    except json.JSONDecodeError as e:
                                        logger.warning(f"解析Dify流式数据失败: {e}, 数据: {data_str}")
                                        continue
                                    except Exception as e:
                                        logger.error(f"处理Dify流式数据时发生错误: {e}")
                                        continue
                                else:
                                    logger.info(f"跳过非data行: {line_str}")
                            else:
                                logger.info("收到空行")
                        
                        logger.info(f"流式数据处理完成，共处理 {line_count} 行")
                        
                    except Exception as e:
                        logger.error(f"流式数据处理过程中发生错误: {e}")
                        import traceback
                        logger.error(f"详细错误信息: {traceback.format_exc()}")
                    
                    # ====== 写入缓存 ======
                    with stream_cache_lock:
                        stream_cache[stream_id] = {"content": stream_content, "finish": True}
                    # 构建最终的流式消息响应
                    if stream_content:
                        stream_message = {
                            "msgtype": "stream",
                            "stream": {
                                "id": stream_id,
                                "finish": True,
                                "content": stream_content
                            }
                        }
                        
                        stream_json_str = json.dumps(stream_message, ensure_ascii=False)
                        logger.info(f"构建的最终流式消息: {stream_message}")
                        
                        # 加密流式消息
                        encrypted_msg, signature = wechat_crypto.encrypt(stream_json_str, timestamp, nonce)
                        
                        if encrypted_msg and signature:
                            # 构建加密回复JSON
                            reply_encrypted_json = {
                                "encrypt": encrypted_msg,
                                "msg_signature": signature,
                                "timestamp": timestamp,
                                "nonce": nonce
                            }
                            
                            logger.info("返回流式消息响应")
                            return jsonify(reply_encrypted_json)
                        else:
                            logger.error("加密流式消息失败")
                            return Response('流式消息加密失败', status=500, content_type='text/plain')
                    else:
                        logger.error("没有获取到Dify回复内容")
                        return Response('没有获取到回复内容', status=500, content_type='text/plain')
                    
                    response = Response(
                        generate_stream(),
                        mimetype='text/event-stream',
                        headers={
                            'Cache-Control': 'no-cache',
                            'Connection': 'keep-alive',
                            'Access-Control-Allow-Origin': '*',
                            'Access-Control-Allow-Headers': 'Cache-Control'
                        }
                    )
                    return response
                else:
                    logger.error("Dify流式API调用失败，使用兼容模式")
                    
                    # 使用兼容模式
                    dify_response = dify_api.send_message(
                        message=user_message,
                        user_id=from_user
                    )
                    
                    logger.info(f"Dify回复: {dify_response}")
                    
                    # 构建企业微信流式消息格式
                    stream_message = {
                        "msgtype": "stream",
                        "stream": {
                            "id": stream_id,
                            "finish": True,
                            "content": dify_response
                        }
                    }
                    
                    stream_json_str = json.dumps(stream_message, ensure_ascii=False)
                    logger.info(f"构建的企业微信流式消息: {stream_message}")
                    
                    # 加密流式消息
                    encrypted_msg, signature = wechat_crypto.encrypt(stream_json_str, timestamp, nonce)
                    
                    if encrypted_msg and signature:
                        # 构建加密回复JSON
                        reply_encrypted_json = {
                            "encrypt": encrypted_msg,
                            "msg_signature": signature,
                            "timestamp": timestamp,
                            "nonce": nonce
                        }
                        
                        logger.info("消息处理完成，返回流式回复")
                        logger.info(f"最终回复JSON: {reply_encrypted_json}")
                        
                        return jsonify(reply_encrypted_json)
                    else:
                        logger.error("消息加密失败")
                        return Response('消息加密失败', status=500, content_type='text/plain')
            else:
                # 非文本消息，返回空回复
                logger.info(f"忽略非文本消息 - 类型: {msg_type}")
                return Response(json.dumps({"status": "success"}), content_type='application/json')
                
        except Exception as e:
            logger.error(f"处理消息时发生错误: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return Response('处理消息失败', status=500, content_type='text/plain')

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    logger.debug("收到健康检查请求")
    return jsonify({
        'status': 'healthy',
        'timestamp': int(time.time())
    })

@app.route('/', methods=['GET'])
def index():
    """首页"""
    logger.debug("收到首页访问请求")
    return jsonify({
        'message': '企业微信机器人服务',
        'version': '1.0.0',
        'status': 'running'
    })

if __name__ == '__main__':
    port = int(os.getenv('PORT', 53230))
    logger.info(f"启动Flask服务器 - 端口: {port}")
    app.run(host='0.0.0.0', port=port, debug=os.getenv('FLASK_DEBUG', 'False').lower() == 'true') 